package com.tuniu.agents.impl.flight;

import com.tuniu.agents.agent.Agent;
import com.tuniu.agents.agent.AgentManagerGroup;
import com.tuniu.agents.common.AgentsConstants;
import com.tuniu.agents.common.agent.SelectionAgent;
import com.tuniu.agents.common.i18n.I18nConstants;
import com.tuniu.agents.common.message.BaseAgentNewResp;
import com.tuniu.agents.common.message.CardTypeEnum;
import com.tuniu.agents.memory.MemoryManager;
import com.tuniu.agents.message.MessageLogger;
import com.tuniu.agents.message.MessageMemoryManager;
import com.tuniu.agents.message.MessageRender;
import com.tuniu.agents.message.MessageSender;
import com.tuniu.agents.tools.flights.constants.FlightResponseConstants;
import com.tuniu.agents.tools.flights.domestic.oneway.FlightInformationQueryTool;
import com.tuniu.agents.tools.flights.domestic.oneway.FlightLowPriceQueryTool;
import com.tuniu.agents.tools.flights.domestic.oneway.FlightTimeQueryTool;
import com.tuniu.agents.tools.flights.domestic.oneway.FlightTransferQueryTool;
import com.tuniu.agents.tools.flights.domestic.roundtrip.RoundFlightInformationQueryTool;
import com.tuniu.agents.tools.flights.domestic.roundtrip.RoundFlightLowPriceQueryTool;
import com.tuniu.agents.tools.flights.international.oneway.IntlFlightInformationQueryTool;
import com.tuniu.agents.tools.flights.international.oneway.IntlFlightLowPriceQueryTool;
import com.tuniu.agents.tools.flights.international.oneway.IntlFlightTransferQueryTool;
import com.tuniu.agents.tools.flights.international.roundtrip.IntlRoundFlightInformationQueryTool;
import com.tuniu.agents.tools.flights.international.roundtrip.IntlRoundFlightLowPriceQueryTool;
import com.tuniu.agents.util.ChatResponseUtils;
import com.tuniu.agents.util.I18nUtil;
import com.tuniu.agents.util.JsonUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;
import reactor.core.publisher.Flux;

import java.util.Collections;
import java.util.List;
import java.util.Map;

import static com.tuniu.agents.impl.flight.FlightAgentConfig.FLIGHT_DOMESTIC_ONE_WAY_ANALYSIS_AGENT_ID;
import static com.tuniu.agents.impl.flight.FlightAgentConfig.FLIGHT_DOMESTIC_ROUND_TRIP_ANALYSIS_AGENT_ID;

@Slf4j
public class FlightTicketAgent extends SelectionAgent {

    private MemoryManager memoryManager;

    private MessageSender messageSender;

    private MessageRender messageRender;

    @Autowired
    private MessageMemoryManager messageMemoryManager;

    public FlightTicketAgent(AgentManagerGroup agentManagerGroup,
                             String agentId, MessageSender messageSender, String promptTemplateId, MessageLogger messageLogger, String... tools) {
        super(agentManagerGroup, agentId, promptTemplateId, messageLogger, null, tools);
        this.messageSender = messageSender;
        this.messageRender = agentManagerGroup.getMessageRender();
        this.memoryManager = agentManagerGroup.getMemoryManager();
    }

    @Override
    public Flux<String> stream(String conversationId, List<Message> messages, Map<String, Object> context) {
        String message = messages.get(0).getText();
        ChatResponse response = this.callResponse(conversationId, message);
        String toolName = ChatResponseUtils.getFirstToolName(response);
        if (StringUtils.isEmpty(toolName)) {
            return Flux.just(ChatResponseUtils.getFirstTextMessage(response));
        } else {
            String toolResponse = ChatResponseUtils.getFirstTextMessage(response);
            CardTypeEnum cradType = getCardTypeByToolName(toolName);
            Agent agent = null;
            String acceptLanguage = I18nUtil.getLocaleLanguageTag(conversationId);
            if (!StringUtils.hasText(toolResponse) || toolResponse.equals("null")) {
                messageMemoryManager.send(conversationId, I18nUtil.getMessage(I18nConstants.COMMON_SYSTEM_ERROR_MESSAGE_KEY), this.agentId, List.of(I18nUtil.getMessage(I18nConstants.COMMON_SYSTEM_ERROR_MESSAGE_KEY)));
                return Flux.just(BaseAgentNewResp.sendAnswerMsg(null, ""));
            }
            switch (toolName) {
                // Domestic one-way flights
                case FlightLowPriceQueryTool.TOOLNAME, FlightTransferQueryTool.TOOLNAME, FlightTimeQueryTool.TOOLNAME:
                    if (CollectionUtils.isEmpty((List<?>) JsonUtil.toMap(toolResponse).get("data"))) {
                        messageMemoryManager.send(conversationId, FlightResponseConstants.COMMON_ONE_FLIGHT_ERROR, this.agentId, List.of(FlightResponseConstants.COMMON_ONE_FLIGHT_ERROR));
                        break;
                    }

                    agent = agentManager.getAgent(FLIGHT_DOMESTIC_ONE_WAY_ANALYSIS_AGENT_ID);
                    break;
                // Domestic round-trip flights
                case RoundFlightLowPriceQueryTool.TOOLNAME:
                    // The mini program currently does not support round-trip queries
                    String clientType = String.valueOf(chatEnvironmentMemory.get(conversationId, AgentsConstants.AgentsConversationEnvironmentKey.CLIENT_TYPE));
                    if (AgentsConstants.FrontendParamKey.WECHAT_CLIENT_TYPE.contains(clientType)) {
                        messageMemoryManager.send(conversationId, FlightResponseConstants.COMMON_FLIGHT_MP_ERROR, this.agentId, List.of(FlightResponseConstants.COMMON_FLIGHT_MP_ERROR));
                        break;
                    }

                    if (CollectionUtils.isEmpty((List<?>) JsonUtil.toMap(toolResponse).get("data"))) {
                        messageMemoryManager.send(conversationId, FlightResponseConstants.COMMON_ROUND_FLIGHT_ERROR, this.agentId, List.of(FlightResponseConstants.COMMON_ROUND_FLIGHT_ERROR));
                        break;
                    }

                    agent = agentManager.getAgent(FLIGHT_DOMESTIC_ROUND_TRIP_ANALYSIS_AGENT_ID);
                    break;
                // One-way flight detail query
                case FlightInformationQueryTool.TOOLNAME, IntlFlightInformationQueryTool.TOOLNAME:
                    Object flightInformation = FlightInformationQueryTool.buildCard(toolResponse, chatEnvironmentMemory.get(conversationId));
                    if (flightInformation == null) {
                        messageMemoryManager.send(conversationId, FlightResponseConstants.COMMON_ONE_FLIGHT_ERROR, this.agentId, List.of(FlightResponseConstants.COMMON_ONE_FLIGHT_ERROR));
                        break;
                    }
                    // Directly return the card
                    List<Object> flightInformationCardInfo = Collections.singletonList(flightInformation);
                    // Flight detail markdown
                    FlightInformationQueryTool.Response markdownResponseOne = JsonUtil.fromJson(toolResponse, FlightInformationQueryTool.Response.class);
                    String markdownFlightOne = messageRender.renderTemplate("flight/flight-info-flight-detail-one.ftlh",Map.of(I18nConstants.LANGUAGE_TAG, acceptLanguage),
                            JsonUtil.toMap(JsonUtil.toJson(markdownResponseOne.flightInfo().get(0))));
                    String markdownCabinOne = messageRender.renderTemplate("flight/flight-info-cabin-detail-one.ftlh",Map.of(I18nConstants.LANGUAGE_TAG, acceptLanguage),
                            JsonUtil.toMap(JsonUtil.toJson(markdownResponseOne)));
                    String flightDetailMarkdownOne = BaseAgentNewResp.sendCardMsg(cradType.getCardType(), flightInformationCardInfo, new StringBuilder().append(markdownFlightOne).append("\n").append(markdownCabinOne).toString(), "");
                    // Simplify card information
                    String flightBaseInfoResult = FlightInformationQueryTool.buildMarkdown((FlightInformationQueryTool.Response) flightInformation);
                    memoryManager.setMessageRight(conversationId, Collections.singletonList(new AssistantMessage(flightBaseInfoResult, Map.of("agentId", this.agentId))));
                    return Flux.just(flightDetailMarkdownOne);
                // Round-trip flight detail query
                case RoundFlightInformationQueryTool.TOOLNAME, IntlRoundFlightInformationQueryTool.TOOLNAME:
                    Object roundFlightCard = RoundFlightInformationQueryTool.buildCard(toolResponse, chatEnvironmentMemory.get(conversationId));
                    if (roundFlightCard == null){
                        messageMemoryManager.send(conversationId, FlightResponseConstants.COMMON_ROUND_FLIGHT_ERROR,this.agentId,List.of(FlightResponseConstants.COMMON_ROUND_FLIGHT_ERROR));
                        break;
                    }
                    // Directly return the card
                    List<Object> roundFlightCardInfo = Collections.singletonList(roundFlightCard);
                    RoundFlightInformationQueryTool.Response markdownResponseRound = JsonUtil.fromJson(toolResponse, RoundFlightInformationQueryTool.Response.class);
                    String markdownFlightRound = messageRender.renderTemplate("flight/flight-info-flight-detail-round.ftlh", Map.of(I18nConstants.LANGUAGE_TAG, acceptLanguage),
                            JsonUtil.toMap(JsonUtil.toJson(markdownResponseRound)));
                    String markdownCabinRound = messageRender.renderTemplate("flight/flight-info-cabin-detail-round.ftlh", Map.of(I18nConstants.LANGUAGE_TAG, acceptLanguage),
                            JsonUtil.toMap(JsonUtil.toJson(markdownResponseRound)));
                    String flightDetailMarkdownRound = BaseAgentNewResp.sendCardMsg(cradType.getCardType(), roundFlightCardInfo, new StringBuilder().append(markdownFlightRound).append("\n").append(markdownCabinRound).toString(), "");
                    // Simplify card information
                    String roundBaseInfoResult = RoundFlightInformationQueryTool.buildMarkdown((RoundFlightInformationQueryTool.Response) roundFlightCard);
                    memoryManager.setMessageRight(conversationId, Collections.singletonList(new AssistantMessage(roundBaseInfoResult, Map.of("agentId", this.agentId))));
                    return Flux.just(flightDetailMarkdownRound);
                default:
                    throw new IllegalArgumentException("Unknown tool name: " + toolName);
            }

            if (agent == null) {
                return Flux.just(BaseAgentNewResp.sendAnswerMsg(null, ""));
            }
            return agent.stream(conversationId, toolResponse);
        }
    }

    private CardTypeEnum getCardTypeByToolName(String toolName) {
        switch (toolName) {
            case FlightLowPriceQueryTool.TOOLNAME, FlightTransferQueryTool.TOOLNAME, FlightTimeQueryTool.TOOLNAME,
                    IntlFlightLowPriceQueryTool.TOOLNAME, IntlFlightTransferQueryTool.TOOLNAME:
                return CardTypeEnum.FLIGHT_LIST_ONE;
            case RoundFlightLowPriceQueryTool.TOOLNAME, IntlRoundFlightLowPriceQueryTool.TOOLNAME:
                return CardTypeEnum.FLIGHT_LIST_ROUND;
            case FlightInformationQueryTool.TOOLNAME, IntlFlightInformationQueryTool.TOOLNAME:
                return CardTypeEnum.FLIGHT_DETAIL;
            case RoundFlightInformationQueryTool.TOOLNAME, IntlRoundFlightInformationQueryTool.TOOLNAME:
                return CardTypeEnum.FLIGHT_ROUND_DETAIL;
            default:
                return CardTypeEnum.FLIGHT_LIST_ONE;
        }
    }
}